import psutil
import json
from typing import Dict, Any

try:
    from mcp import ClientSession, StdioServerParameters
    from mcp.client.stdio import stdio_client
    from mcp.server.fastmcp import FastMCP
    MCP_AVAILABLE = True
except ImportError:
    print("警告: 未安装MCP库。请运行 'pip install mcp' 安装。")
    ClientSession = None
    StdioServerParameters = None
    stdio_client = None
    FastMCP = None
    MCP_AVAILABLE = False

def get_cpu_usage() -> Dict[str, Any]:
    """
    获取CPU使用率信息
    
    Returns:
        包含CPU使用率信息的字典
    """
    cpu_percent = psutil.cpu_percent(interval=1)
    cpu_count = psutil.cpu_count()
    cpu_count_logical = psutil.cpu_count(logical=True)
    
    return {
        "cpu_percent": cpu_percent,
        "cpu_count": cpu_count,
        "cpu_count_logical": cpu_count_logical
    }

def get_memory_usage() -> Dict[str, Any]:
    """
    获取内存使用率信息
    
    Returns:
        包含内存使用率信息的字典
    """
    memory = psutil.virtual_memory()
    
    return {
        "total_gb": round(memory.total / (1024**3), 2),
        "available_gb": round(memory.available / (1024**3), 2),
        "used_gb": round(memory.used / (1024**3), 2),
        "percent": memory.percent
    }

def get_disk_usage(path: str = "/") -> Dict[str, Any]:
    """
    获取磁盘使用率信息
    
    Args:
        path: 磁盘路径，默认为根路径
        
    Returns:
        包含磁盘使用率信息的字典
    """
    try:
        disk = psutil.disk_usage(path)
        return {
            "total_gb": round(disk.total / (1024**3), 2),
            "used_gb": round(disk.used / (1024**3), 2),
            "free_gb": round(disk.free / (1024**3), 2),
            "percent": round(disk.used / disk.total * 100, 2)
        }
    except FileNotFoundError:
        # Windows系统可能需要使用其他路径
        disk = psutil.disk_usage("C:\\")
        return {
            "total_gb": round(disk.total / (1024**3), 2),
            "used_gb": round(disk.used / (1024**3), 2),
            "free_gb": round(disk.free / (1024**3), 2),
            "percent": round(disk.used / disk.total * 100, 2)
        }

def get_system_info() -> Dict[str, Any]:
    """
    获取完整的系统信息
    
    Returns:
        包含所有系统信息的字典
    """
    return {
        "cpu": get_cpu_usage(),
        "memory": get_memory_usage(),
        "disk": get_disk_usage()
    }

# 如果安装了MCP库，则提供MCP工具封装
if MCP_AVAILABLE:
    # 创建FastMCP实例
    mcp = FastMCP("System Resource Monitor")
    
    @mcp.tool()
    def check_cpu_usage() -> Dict[str, Any]:
        """
        检查CPU使用率 (MCP工具版本)
        
        Returns:
            包含CPU使用率信息的字典
        """
        return get_cpu_usage()
    
    @mcp.tool()
    def check_memory_usage() -> Dict[str, Any]:
        """
        检查内存使用率 (MCP工具版本)
        
        Returns:
            包含内存使用率信息的字典
        """
        return get_memory_usage()
    
    @mcp.tool()
    def check_disk_usage(path: str = "/") -> Dict[str, Any]:
        """
        检查磁盘使用率 (MCP工具版本)
        
        Args:
            path: 磁盘路径，默认为根路径
            
        Returns:
            包含磁盘使用率信息的字典
        """
        return get_disk_usage(path)
    
    @mcp.tool()
    def check_system_status() -> Dict[str, Any]:
        """
        检查完整系统状态 (MCP工具版本)
        
        Returns:
            包含所有系统信息的字典
        """
        return get_system_info()

def main():
    """
    主函数，用于命令行测试
    """
    print("系统资源监控工具")
    print("=" * 30)
    
    # 获取并显示CPU信息
    cpu_info = get_cpu_usage()
    print(f"CPU使用率: {cpu_info['cpu_percent']}%")
    print(f"CPU核心数: {cpu_info['cpu_count']} 物理核心, {cpu_info['cpu_count_logical']} 逻辑核心")
    
    print("\n" + "-" * 30)
    
    # 获取并显示内存信息
    memory_info = get_memory_usage()
    print(f"内存总量: {memory_info['total_gb']} GB")
    print(f"已用内存: {memory_info['used_gb']} GB")
    print(f"可用内存: {memory_info['available_gb']} GB")
    print(f"内存使用率: {memory_info['percent']}%")
    
    print("\n" + "-" * 30)
    
    # 获取并显示磁盘信息
    disk_info = get_disk_usage()
    print(f"磁盘总量: {disk_info['total_gb']} GB")
    print(f"已用磁盘: {disk_info['used_gb']} GB")
    print(f"可用磁盘: {disk_info['free_gb']} GB")
    print(f"磁盘使用率: {disk_info['percent']}%")
    
    # 如果有MCP库，提供启动服务器选项
    if MCP_AVAILABLE:
        print("\n" + "=" * 30)
        print("MCP服务器可用，可以作为工具提供给其他应用使用")

if __name__ == "__main__":
    main()